import pandas as pd
import numpy as np
from training.ml_frame.step_rolling.ml_train import Ml_rolling_trin
import time

class Rolling_train_twoset():
    
    def __init__(self, feature, label, codes, times, args) -> None:
        self.args = args
        self.feature = feature
        self.label = label 
        self.codes = codes
        self.times = times
        
    def rolling_train(self):
        test_lst = []
        out_id =  np.where(self.times>=pd.to_datetime(self.args.time_param['outsample_beg']))[0]
        out_length = len(out_id)
        out_beg_id = out_id[0]
        mrt = Ml_rolling_trin(self.args)
        for i in range(0, out_length, self.args.rolling_step):
            t1 = time.time()
            train_x = self.feature[:out_beg_id+i-self.args.label_range,:]
            train_y = self.label[:out_beg_id+i-self.args.label_range,:]
            test_x = self.feature[out_beg_id+i:out_beg_id+i+self.args.rolling_step,:]
            test_y = self.label[out_beg_id+i:out_beg_id+i+self.args.rolling_step,:]         
            train_pred, test_pred = mrt.forward_train(train_x, test_x, train_y, test_y)#两种训练方式，保留图像与IC 
            test_lst.append(test_pred)   
            print(f'滚动length：{i}/{out_length}, 耗时：{(time.time()-t1):.2f}s') 
        test_ary = np.concatenate(test_lst, axis=0)
        train_times = self.times[out_beg_id:]
        test_times = self.times[:out_beg_id]
        test_pred_df = pd.DataFrame(test_ary, index=self.times[out_beg_id:], columns=self.codes)
        test_true_df = pd.DataFrame(self.label[out_beg_id:].squeeze(-1), index=self.times[out_beg_id:], columns=self.codes)
        train_pred_df = pd.DataFrame(train_pred, index=self.times[:out_beg_id+i-self.args.label_range], columns=self.codes)
        train_true_df = pd.DataFrame(self.label[:out_beg_id+i-self.args.label_range].squeeze(-1), index=self.times[:out_beg_id+i-self.args.label_range], columns=self.codes)
        return train_pred_df, train_true_df, test_pred_df, test_true_df,  train_times, test_times